/*
 * KiKernel.cpp
 *
 *  Created on: 2013-10-18
 *      Author: WinDDK
 */

#include "KiKernel.h"

KiKernel::KiKernel()
{
    mHook = NULL;
    mWin32kPE = NULL;
    //mDebug = NULL;
    KiDebug("Success");

}

KiKernel::~KiKernel()
{
    if(mHook != NULL)
    {
        delete mHook;
    }

    if(mWin32kPE != NULL)
    {
        delete mWin32kPE;
    }

//    if(mDebug != NULL)
//    {
//       delete mDebug;
//    }
    KiDebug("Success");
}

BOOLEAN KiKernel::GetModulePath()
{
    strcpy(mModulePath, "\\??\\");
    if (!GetSystemPath(this->mSystemPath))
    {
        KiDebug("GetSystemPath Failed");
        return FALSE;
    }
    if (!GetKernelName(this->mModuleName))
    {
        KiDebug("GetKernelModuleName Failed");
        return FALSE;
    }
    strcat(mModulePath, mSystemPath);
    strcat(mModulePath, "\\");
    strcat(mModulePath, mModuleName);
    KiDebug("%s", mModulePath);
    return TRUE;
}

// 得到系统内核模块文件路径
BOOLEAN KiKernel::GetKernelName(OUT PCHAR moduleName)
{
    ULONG size;
    PSYSTEM_MODULES modules = NULL;
    ULONG status = ZwQuerySystemInformation(SystemModuleInformation, NULL, 0, &size);
    if (status == STATUS_INFO_LENGTH_MISMATCH )
    {
        modules = (PSYSTEM_MODULES) ExAllocatePoolWithTag(PagedPool, size, 'KiPE');
        status = ZwQuerySystemInformation(SystemModuleInformation, modules, size, NULL);
        if (!NT_SUCCESS(status) || modules == NULL)
        {
            KiDebug("ZwQuerySystemInformation Failed");
            ExFreePool((PVOID ) modules);
            return FALSE;
        }
    } else
    {
        KiDebug("ZwQuerySystemInformation Failed");
        return FALSE;
    }
    strncpy(moduleName, modules->Modules[0].ImageName + modules->Modules[0].ModuleNameOffset, MODULENAMELENGTH);
    KiDebug("%s", moduleName);
    return TRUE;
}

BOOLEAN KiKernel::LoadWin32K()
{
    mWin32kPE = new (NonPagedPool) KiWin32K();
    if(mWin32kPE == NULL)
    {
        return FALSE;
    }

    if (!mWin32kPE->SetKernelInfo(mModuleName, mNewImageBase))
    {
        KiDebug("mWin32kPE->SetKernelInfo Failed");
        return FALSE;
    }

    if (!mWin32kPE->LoadModule())
    {
        KiDebug("mWin32kPE->LoadModule Failed");
        return FALSE;
    }

    if (!mWin32kPE->FixModule())
    {
        KiDebug("mWin32kPE->FixModule Failed");
        return FALSE;
    }
    return TRUE;
}

BOOLEAN KiKernel::HookKiFastCallEntry()
{
    ULONG oldSSDT = (ULONG) (KeServiceDescriptorTable);
    ULONG newSSDT = (ULONG) (oldSSDT) + mNewImageBase - mOldImageBase;
    PSYSTEM_SERVICE_TABLE oldShadowSSDT;
    if(!GetOldShadowSSDT(oldShadowSSDT))
    {
        KiDebug("GetOldShadowSSDT Failed");
        return FALSE;
    }
    oldShadowSSDT++;
    ULONG newShadowSSDT = (ULONG) (oldShadowSSDT) + mNewImageBase - mOldImageBase;
    mHook = new (NonPagedPool) KiHook(oldSSDT, newSSDT, (ULONG) (oldShadowSSDT), newShadowSSDT);
    if(mHook == NULL)
    {
        return FALSE;
    }
    return TRUE;
}

BOOLEAN KiKernel::FixDebugPort()
{
//    mDebug = new (NonPagedPool)KiDebugPort(mOldImageBase, mNewImageBase);
//
//    if (!mDebug->FixDebugPortCode())
//    {
//        KiDebug("debugport.FixDebugPortCode Failed");
//        return FALSE;
//    }
//
//    if (!mDebug->FixCodeReferences())
//    {
//        KiDebug("debugport.FixCodeReferences");
//        return FALSE;
//    }
//
//    if (!mDebug->FixInterruptTable())
//    {
//        KiDebug("debugport.FixInterruptTable");
//        return FALSE;
//    }

    return TRUE;
}

BOOLEAN KiKernel::FixModule()
{
    if (!FixRelocTable())
    {
        KiDebug("FixRelocTable Failed");
        return FALSE;
    }

    if (!FixImportTable())
    {
        KiDebug("FixImportTable Failed");
        return FALSE;
    }

    if (!FixSSDT(KeServiceDescriptorTable, (PSYSTEM_SERVICE_TABLE)((ULONG)KeServiceDescriptorTable + mNewImageBase - mOldImageBase), mNewImageBase, mOldImageBase))
    {
        KiDebug("FixSSDT Failed");
        return FALSE;
    }

    if(!LoadWin32K())
    {
        KiDebug("LoadWin32K Failed");
        return FALSE;
    }

    if(!FixShadowSSDT())
    {
        KiDebug("FixShadowSSDT Failed");
        return FALSE;
    }

    if(!FixDebugPort())
    {
        KiDebug("FixDebugPort Failed");
        return FALSE;
    }

    if(!HookKiFastCallEntry())
    {
        KiDebug("HookKiFastCallEntry Failed");
        return FALSE;
    }

    return TRUE;
}

BOOLEAN KiKernel::GetOldShadowSSDT(PSYSTEM_SERVICE_TABLE & oldShadowSSDT)
{
    PUCHAR ptr = (PUCHAR) KeAddSystemServiceTable;
    for (ULONG i = 0; i < 1024; i++)
    {
        ULONG address = *(PLONG) (ptr + i);
        if (MmIsAddressValid((PVOID) address) && MmIsAddressValid((PVOID) (address + sizeof(SYSTEM_SERVICE_TABLE))))
        {
            if (memcmp((PVOID)address, KeServiceDescriptorTable, sizeof(SYSTEM_SERVICE_TABLE)) == 0)
            {
                if ( address == (ULONG)KeServiceDescriptorTable)
                {
                    continue;
                }
                oldShadowSSDT = (PSYSTEM_SERVICE_TABLE) address;
                KiDebug("Old ShadowSSDT address is 0x%x", (ULONG )oldShadowSSDT);
                return TRUE;
            }
        }
    }
    return FALSE;
}

BOOLEAN KiKernel::FixSSDT(PSYSTEM_SERVICE_TABLE oldSSDT, PSYSTEM_SERVICE_TABLE newSSDT, ULONG newImageBase, ULONG oldImageBase)
{

    KiDebug("Old SSDT address is 0x%x", (ULONG)oldSSDT);
    KiDebug("New SSDT address is 0x%x", (ULONG)newSSDT);
    if (!MmIsAddressValid(oldSSDT) || !MmIsAddressValid(newSSDT) || newImageBase == oldImageBase)
    {
        KiDebug("newSSDT is unaviable!");
        return FALSE;
    }

    newSSDT->ServiceTable = NULL;
    if (oldSSDT->ServiceTable != NULL)
    {
        newSSDT->ServiceTable = (PULONG) ((ULONG) oldSSDT->ServiceTable + newImageBase - oldImageBase);
        if (!MmIsAddressValid(newSSDT->ServiceTable))
        {
            KiDebug("newSSDT->ServiceTable: 0x%x", newSSDT->ServiceTable);
            return FALSE;
        }
        for (ULONG i = 0; i < oldSSDT->NumberOfService; i++)
        {
             newSSDT->ServiceTable[i] = oldSSDT->ServiceTable[i] + newImageBase - oldImageBase;
        }
    }

    newSSDT->ServiceCounterTable = NULL;
    if (oldSSDT->ServiceCounterTable != NULL)
    {
        newSSDT->ServiceCounterTable = (PULONG) ((ULONG) oldSSDT->ServiceCounterTable + newImageBase - oldImageBase);
        if (!MmIsAddressValid(newSSDT->ServiceCounterTable))
        {
            KiDebug("newSSDT->ServiceCounterTable: 0x%x", newSSDT->ServiceCounterTable);
            return FALSE;
        }
        *(newSSDT->ServiceCounterTable) = *(oldSSDT->ServiceCounterTable);
    }

    newSSDT->NumberOfService = oldSSDT->NumberOfService;

    newSSDT->ParamTable = NULL;
    if (oldSSDT->ParamTable != NULL)
    {
        newSSDT->ParamTable = (PUCHAR) ((ULONG) oldSSDT->ParamTable + newImageBase - oldImageBase);
        if (!MmIsAddressValid(newSSDT->ParamTable))
        {
            KiDebug("newSSDT->ParamTable: 0x%x", newSSDT->ParamTable);
            return FALSE;
        }
        RtlCopyMemory(newSSDT->ParamTable, oldSSDT->ParamTable, oldSSDT->NumberOfService);
    }
    return TRUE;
}

BOOLEAN KiKernel::FixShadowSSDT()
{
    PSYSTEM_SERVICE_TABLE oldSSDT = KeServiceDescriptorTable;
    PSYSTEM_SERVICE_TABLE newSSDT = (PSYSTEM_SERVICE_TABLE) ((ULONG) oldSSDT + mNewImageBase - mOldImageBase);

    PSYSTEM_SERVICE_TABLE oldShadowSSDT;
    if (!GetOldShadowSSDT(oldShadowSSDT))
    {
        KiDebug("GetOldShadowSSDT Failed");
        return FALSE;
    }
    PSYSTEM_SERVICE_TABLE newShadowSSDT = (PSYSTEM_SERVICE_TABLE) ((ULONG) oldShadowSSDT + mNewImageBase - mOldImageBase);
    KiDebug(" New ShadowSSDT address is 0x%x", (ULONG )newShadowSSDT);
    RtlCopyMemory(newShadowSSDT, newSSDT, sizeof(SYSTEM_SERVICE_TABLE));

    oldShadowSSDT++;
    newShadowSSDT++;

    ULONG CSRSSPID;
    PEPROCESS crsEProc;

    if(!GetPIDByProcessName(L"csrss.exe", CSRSSPID))
    {
        KiDebug("GetPIDByProcessName Failed");
        return FALSE;
    }
    NTSTATUS status = PsLookupProcessByProcessId(CSRSSPID, &crsEProc);
    if (!NT_SUCCESS(status))
    {
        KiDebug("PsLookupProcessByProcessId Failed");
        return FALSE;
    }

    status = KeAttachProcess(crsEProc);

    if(!FixSSDT(oldShadowSSDT, newShadowSSDT, mWin32kPE->GetNewImageBase(), mWin32kPE->GetOldImageBase()))
    {
        return FALSE;
    }

    status = KeDetachProcess();

    return TRUE;
}

BOOLEAN KiKernel::GetPIDByProcessName(PWCHAR processName, ULONG & processId)
{
    ULONG size;
    NTSTATUS status = ZwQuerySystemInformation(SystemProcessInformation, NULL, 0, &size);
    if (status != STATUS_INFO_LENGTH_MISMATCH )
    {
        KiDebug("ZwQuerySystemInformation Failed");
        return FALSE;
    }

    PVOID buffer = ExAllocatePoolWithTag(NonPagedPool, size, KIKERNEL);
    PSYSTEM_PROCESS_INFORMATION processes = (PSYSTEM_PROCESS_INFORMATION) buffer;
    if (processes)
    {
        status = ZwQuerySystemInformation(SystemProcessInformation, processes, size, &size);
        if (!NT_SUCCESS(status))
        {
            KiDebug("ZwQuerySystemInformation failed.");
            ExFreePool((PVOID ) buffer);
            return FALSE;
        }
    } else
    {
        KiDebug("ExAllocatePool error\n");
        return FALSE;
    }

    while (1)
    {
        if (processes->ProcessName.Buffer != NULL && !wcscmp(processes->ProcessName.Buffer, processName))
        {
            KiDebug("%S Processes's id is: %d", processes->ProcessName.Buffer, processes->ProcessId);
            processId = processes->ProcessId;
            ExFreePool(buffer);
            return TRUE;
        }
        if (processes->NextEntryDelta == 0)
        {
            break;
        }
        processes = (PSYSTEM_PROCESS_INFORMATION) (((PUCHAR) processes) + processes->NextEntryDelta);
    }

    ExFreePool(buffer);
    return FALSE;
}
